import os
import argparse
import json
from .net import model_dict


def start_application(model_name, model_path, data_path, save_path, data):
    os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
    save_path = os.path.join(save_path, model_name)

    os.makedirs(save_path, exist_ok=True)

    model_config_path = os.path.join(model_path, "model_config.json")
    meta_information_path = os.path.join(data_path, "meta_information.json")

    print("\n\n # Loading the Configurations :")

    if os.path.isfile(meta_information_path):
        with open(meta_information_path, "r") as file:
            meta_information = json.load(file)
    else:
        raise ValueError("no data")

    if os.path.isfile(model_config_path):
        with open(model_config_path, "r") as file:
            model_config = json.load(file)
    else:
        raise ValueError("no model")

    print("# Used model_config: ", model_config)

    try:
        if model_config['default'] == True:
            print(
                '\n\t# WARNING : YOU ARE USING THE DEFAULT model_config.json GENERATED DURING DATA PREPARATION, YOU CAN MODIFY IT AT :',
                model_config_path)
    except:
        pass

    model_cls = model_dict[model_name]
    model = model_cls(data_path, model_config_path, save_path,
                      meta_information, model_config, data=data, train_config=None)

    y_pred = model.predict()
    return y_pred


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument("--data-path", type=str, dest="data_path", default="./data/system_csv/",
                        help="set the direction path of ptm dataset")
    parser.add_argument("--model-name", type=str, dest="model_name",
                        choices=["convTransformer", "seq2seq", "pyraformer", "informer", "lstm", "gru", "transformer",
                                 "crate", "mtct"],
                        help="set the model to be trained, model name must be in one of convTransformer, seq2seq, pyraformer, informer, lstm, gru, crate")
    parser.add_argument("--model-path", type=str, dest="model_path", default="./model",
                        help="set the path of model to be loaded")
    parser.add_argument("--save-path", type=str, dest="save_path", default="./prediction",
                        help="set the saved model direcation path")
    parser.set_defaults(prep=True)

    return parser.parse_args()

def start(data_dir = None):
    # if __name__ == "__main__":
    args = parse_args()
    model_name = args.model_name
    data_path = args.data_path
    save_path = args.save_path
    model_path = args.model_path

    model_path = os.path.join(model_path, model_name)

    start_application(model_name, model_path, data_path, save_path)

if __name__ == "__main__":

    start()
    # args = parse_args()
    # model_name = args.model_name
    # data_path = args.data_path
    # save_path = args.save_path
    # model_path = args.model_path
    #
    # model_path = os.path.join(model_path, model_name)
    #
    # main(model_name, model_path, data_path, save_path)